package regression

import (
	"fmt"
	"math"

	"gonum.org/v1/gonum/mat"
	"gonum.org/v1/gonum/stat"
)

type RegressionResult struct {
	Intercept    float64
	Coefficients []float64
}

func sumSq(v []float64) float64 {
	var sum float64
	for _, val := range v {
		sum += val * val
	}
	return sum
}

type core struct {
	matX      *mat.Dense
	vecY      []float64
	matReg    *mat.Dense
	coeffs    []float64
	residuals []float64
	diagXXInv []float64
}

type Residual struct {
	Predictions []float64
	Residuals   []float64
}

func (residual *Residual) compute(c *core) error {
	n, pPlus1 := c.matX.Dims()

	betaVec := mat.NewVecDense(pPlus1, c.coeffs)
	yPredVec := mat.NewVecDense(n, nil)
	yPredVec.MulVec(c.matX, betaVec)
	yPred := make([]float64, n)
	for i := range yPred {
		yPred[i] = yPredVec.At(i, 0)
	}

	residuals := make([]float64, n)
	for i := range residuals {
		residuals[i] = c.vecY[i] - yPred[i]
	}
	residual.Predictions = yPred
	residual.Residuals = residuals

	c.residuals = residuals
	return nil
}

type Statistic struct {
	StdErrors   []float64
	TValues     []float64
	RSquared    float64
	AdjustedRSq float64
}

func (statistic *Statistic) compute(c *core) error {
	n, pPlus1 := c.matX.Dims()
	p := pPlus1 - 1

	SSres := sumSq(c.residuals)
	yMean := stat.Mean(c.vecY, nil)
	SStot := 0.0
	for _, yi := range c.vecY {
		SStot += (yi - yMean) * (yi - yMean)
	}

	rSquared := 1 - SSres/SStot
	dfTotal := float64(n - 1)
	dfResidual := float64(n - p - 1)
	adjustedRSq := 1 - (SSres/dfResidual)/(SStot/dfTotal)

	s2 := SSres / dfResidual

	invR := mat.NewTriDense(pPlus1, true, nil)
	invR.Copy(c.matReg)
	if err := invR.InverseTri(invR); err != nil {
		return fmt.Errorf("无法求逆 R 矩阵: %v", err)
	}

	diagXXInv := make([]float64, pPlus1)
	for i := 0; i < pPlus1; i++ {
		sum := 0.0
		for j := i; j < pPlus1; j++ {
			val := invR.At(i, j)
			sum += val * val
		}
		diagXXInv[i] = sum
	}

	stdErrors := make([]float64, pPlus1)
	for i := 0; i < pPlus1; i++ {
		stdErrors[i] = math.Sqrt(s2 * diagXXInv[i])
	}

	tValues := make([]float64, pPlus1)
	for i := 0; i < pPlus1; i++ {
		if stdErrors[i] == 0 {
			tValues[i] = math.Inf(1)
		} else {
			tValues[i] = c.coeffs[i] / stdErrors[i]
		}
	}

	statistic.StdErrors = stdErrors
	statistic.TValues = tValues
	statistic.RSquared = rSquared
	statistic.AdjustedRSq = adjustedRSq
	c.diagXXInv = diagXXInv
	return nil
}

type VIF struct {
	VIFs []float64
}

func (vif *VIF) compute(c *core) error {
	n, pPlus1 := c.matX.Dims()
	p := pPlus1 - 1

	// === 优化后的 VIF 计算（替代辅助回归）===
	vifs := make([]float64, p)
	for j := 0; j < p; j++ {
		colIndex := j + 1 // 跳过截距列（第0列）

		// 提取预测变量列
		col := make([]float64, n)
		for i := 0; i < n; i++ {
			col[i] = c.matX.At(i, colIndex)
		}

		// 计算中心化后的总平方和（TSS）
		mean := stat.Mean(col, nil)
		tss := 0.0
		for _, val := range col {
			diff := val - mean
			tss += diff * diff
		}

		// 处理常数预测变量
		if tss < 1e-15 {
			vifs[j] = math.Inf(1)
		} else {
			vifs[j] = c.diagXXInv[colIndex] * tss
		}
	}
	vif.VIFs = vifs
	return nil
}

type regression struct {
	x [][]float64
	y []float64
}

func New(X [][]float64, y []float64) *regression {
	this := &regression{
		x: X,
		y: y,
	}
	return this
}
func (rb *regression) Run() (*RegressionResult, error) {
	x, y, err := makeMatrix(rb.x, rb.y)
	if err != nil {
		return nil, fmt.Errorf("makeMatrix failed, %v", err)
	}
	result, _, err := coreRegression(x, y)
	if err != nil {
		return nil, fmt.Errorf("coreRegression failed, %v", err)
	}
	return result, nil
}

func (rb *regression) SetResidual(r *Residual) *regression2 {
	this := &regression2{
		x: rb.x,
		y: rb.y,
		r: r,
	}
	return this
}

type regression2 struct {
	x [][]float64
	y []float64
	r *Residual
}

func (rb2 *regression2) Run() (*RegressionResult, error) {
	x, y, err := makeMatrix(rb2.x, rb2.y)
	if err != nil {
		return nil, fmt.Errorf("makeMatrix failed, %v", err)
	}
	result, c, err := coreRegression(x, y)
	if err != nil {
		return nil, fmt.Errorf("coreRegression failed, %v", err)
	}
	err = rb2.r.compute(c)
	if err != nil {
		return nil, fmt.Errorf("r.compute failed, %v", err)
	}
	return result, nil
}

func (rb2 *regression2) SetStatistic(s *Statistic) *regression3 {
	this := &regression3{
		x: rb2.x,
		y: rb2.y,
		r: rb2.r,
		s: s,
	}
	return this
}

type regression3 struct {
	x [][]float64
	y []float64
	r *Residual
	s *Statistic
}

func (rb3 *regression3) Run() (*RegressionResult, error) {
	x, y, err := makeMatrix(rb3.x, rb3.y)
	if err != nil {
		return nil, fmt.Errorf("makeMatrix failed, %v", err)
	}
	result, c, err := coreRegression(x, y)
	if err != nil {
		return nil, fmt.Errorf("coreRegression failed, %v", err)
	}
	err = rb3.r.compute(c)
	if err != nil {
		return nil, fmt.Errorf("r.compute failed, %v", err)
	}
	err = rb3.s.compute(c)
	if err != nil {
		return nil, fmt.Errorf("s.compute failed, %v", err)
	}
	return result, nil
}

func (rb3 *regression3) SetVIF(v *VIF) *regression4 {
	this := &regression4{
		x: rb3.x,
		y: rb3.y,
		r: rb3.r,
		s: rb3.s,
		v: v,
	}
	return this
}

type regression4 struct {
	x [][]float64
	y []float64
	r *Residual
	s *Statistic
	v *VIF
}

func (rb4 *regression4) Run() (*RegressionResult, error) {
	x, y, err := makeMatrix(rb4.x, rb4.y)
	if err != nil {
		return nil, fmt.Errorf("makeMatrix failed, %v", err)
	}
	result, c, err := coreRegression(x, y)
	if err != nil {
		return nil, fmt.Errorf("coreRegression failed, %v", err)
	}
	err = rb4.r.compute(c)
	if err != nil {
		return nil, fmt.Errorf("r.compute failed, %v", err)
	}
	err = rb4.s.compute(c)
	if err != nil {
		return nil, fmt.Errorf("s.compute failed, %v", err)
	}
	err = rb4.v.compute(c)
	if err != nil {
		return nil, fmt.Errorf("v.compute failed, %v", err)
	}
	return result, nil
}

// 优化数据拷贝：按列填充矩阵（提升缓存效率）
func makeMatrix(X [][]float64, y []float64) (*mat.Dense, *mat.Dense, error) {
	if len(y) == 0 {
		return nil, nil, fmt.Errorf("y 向量不能为空")
	}
	n := len(y)

	if len(X) == 0 {
		return nil, nil, fmt.Errorf("X 矩阵不能为空")
	}

	isTransposed := len(X[0]) == n
	p := len(X)
	if !isTransposed {
		if len(X) != n {
			return nil, nil, fmt.Errorf("X 矩阵维度与 y 向量长度不匹配")
		}
		if n > 0 {
			p = len(X[0])
		}
	}

	XMatrix := mat.NewDense(n, p+1, nil)
	YMatrix := mat.NewDense(n, 1, y)

	// 设置截距列（第0列）
	for i := 0; i < n; i++ {
		XMatrix.Set(i, 0, 1)
	}

	// 按列填充（匹配 Gonum 的列主序存储）
	if isTransposed {
		for j := 0; j < p; j++ { // 对每个预测变量
			for i := 0; i < n; i++ {
				XMatrix.Set(i, j+1, X[j][i])
			}
		}
	} else {
		for j := 0; j < p; j++ { // 对每个预测变量
			for i := 0; i < n; i++ {
				XMatrix.Set(i, j+1, X[i][j])
			}
		}
	}
	return XMatrix, YMatrix, nil
}

// 优化后的核心回归函数：使用直接公式计算 VIF
func coreRegression(X *mat.Dense, y *mat.Dense) (*RegressionResult, *core, error) {
	n, pPlus1 := X.Dims()
	yRows, yCols := y.Dims()
	if yRows != n || yCols != 1 {
		return nil, nil, fmt.Errorf("y 必须是 %dx1 矩阵，实际为 %dx%d", n, yRows, yCols)
	}

	yVec := make([]float64, n)
	for i := 0; i < n; i++ {
		yVec[i] = y.At(i, 0)
	}

	qr := new(mat.QR)
	qr.Factorize(X)
	q := new(mat.Dense)
	reg := new(mat.Dense)
	qr.QTo(q)
	qr.RTo(reg)

	qtr := q.T()
	qty := new(mat.Dense)
	qty.Mul(qtr, y)

	coeffs := make([]float64, pPlus1)
	for i := pPlus1 - 1; i >= 0; i-- {
		if math.Abs(reg.At(i, i)) < 1e-15 {
			return nil, nil, fmt.Errorf("矩阵在位置 (%d,%d) 奇异", i, i)
		}
		coeffs[i] = qty.At(i, 0)
		for j := i + 1; j < pPlus1; j++ {
			coeffs[i] -= coeffs[j] * reg.At(i, j)
		}
		coeffs[i] /= reg.At(i, i)
	}

	c := &core{
		matX:   X,
		vecY:   yVec,
		coeffs: coeffs,
		matReg: reg,
	}

	return &RegressionResult{
		Intercept:    coeffs[0],
		Coefficients: coeffs[1:],
	}, c, nil
}
